
import helper
import numpy as np
import rasterio as rio

# 一般地，3 bands 视为 RGB，4 bands 视为 RGB+NIR
def main(args):
    if args.dst is None or args.dst == args.src:
        raise ValueError('dst must be specified and different from src')

    with rio.io.MemoryFile() as memfile:

        with rio.open(args.src, 'r') as src:
            kwargs = src.profile.copy()
            if args.normalize:
                kwargs.update(
                    dtype=rio.uint8,
                )
            if args.reproject and src.crs is None:
                # 地理围栏（左、下、右、上，视为西、南、东、北）
                bounds = rio.coords.BoundingBox(0.0, 0.0, src.width, src.height)
                # https://epsg.io/3857
                kwargs.update(
                    # crs=rio.CRS.from_epsg(3857),
                    transform=rio.transform.from_bounds(*bounds, width=src.width, height=src.height),
                )
            # 只能写打开一次
            with memfile.open(**kwargs) as tmp:
                arr = src.read()
                # (bands, height, width)
                print(arr.shape)

                if args.normalize:
                    arr = helper.remap_to_rgb(arr, lp=args.lp, hp=args.hp)

                if args.bgr:
                    arr = arr[(2, 1, 0), :, :]

                tmp.write(arr)

        with memfile.open() as tmp:
            with rio.open(args.dst, 'w', **tmp.meta) as dst:
                arr = tmp.read()
                dst.write(arr)

    with rio.open(args.dst, 'r+') as dst:

        if args.set_color is not None:
            if src.count == 1:
                dst.colorinterp = (rio.enums.ColorInterp.gray,)
            if src.count == 3:
                if args.set_color == 'gray':
                    dst.colorinterp = (
                        rio.enums.ColorInterp.gray,
                        rio.enums.ColorInterp.undefined,
                        rio.enums.ColorInterp.undefined,
                    )
                else:
                    dst.colorinterp = (
                        rio.enums.ColorInterp.red,
                        rio.enums.ColorInterp.green,
                        rio.enums.ColorInterp.blue,
                    )
            if src.count == 4:
                if args.set_color == 'gray':
                    dst.colorinterp = (
                        rio.enums.ColorInterp.gray,
                        rio.enums.ColorInterp.undefined,
                        rio.enums.ColorInterp.undefined,
                        rio.enums.ColorInterp.undefined,
                    )
                else:
                    dst.colorinterp = (
                        rio.enums.ColorInterp.red,
                        rio.enums.ColorInterp.green,
                        rio.enums.ColorInterp.blue,
                        rio.enums.ColorInterp.undefined,
                    )


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src')
    parser.add_argument('dst')
    parser.add_argument('--lp', default=0.10, type=float, help='lower percent')
    parser.add_argument('--hp', default=99.90, type=float, help='higher percent')
    parser.add_argument('--normalize', action='store_true', help='normalize')
    parser.add_argument('--set-color', choices=['gray', 'rgb'])
    parser.add_argument('--bgr', action='store_true')
    parser.add_argument('--reproject', action='store_true')
    args = parser.parse_args()
    main(args)
